import copy
import os
from datetime import datetime

import matplotlib.pyplot as plt
import torch
from torch.utils import data as Data
from model import GoogLetNetModel as ModelNet
from torch import nn
import time
import pandas as pd
from dataset import GetDataLoader

# 创建存放训练完成后模型的路径
def create_model_save_path():
    model_path = os.path.dirname(ModelNet.PTH_PATH)
    if not os.path.exists(model_path):
        os.makedirs(model_path)
    return model_path

# 模型训练的流程
def train_model_process(save_path:str,model:nn.Module,train_dataloader:Data.DataLoader,val_dataloader:Data.DataLoader,num_epochs:int):
    '''
    :param save_path: 训练好的模型的参数保存的路径
    :param model: 模型对象
    :param train_dataloader:训练加载的数据
    :param val_dataloader: 测试加载的数据
    :param num_epochs: 训练的轮次
    :return:
    '''
    # 使用gpu或cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 优化器：(优化器是在梯度下降法上进行优化后的)
    optimizer = torch.optim.Adam(model.parameters(),lr=0.001)

    # 损失函数: 当前是交叉熵损失函数（分类使用）【回归使用的是均方差】，利用损失值来更新参数的
    criterion = nn.CrossEntropyLoss()

    # 将模型放到训练设备中
    model = model.to(device)

    # 复制当前模型的参数：
    ## 初始化时是有初始参数，后续训练时，更新参数后，保存更新后的参数
    best_model_wts = copy.deepcopy(model.state_dict())

    ## 初始化参数，
    # 最高准确度
    best_acc = 0.0
    # 训练集损失值的列表
    train_loss_all = []
    # 验证集损失值的列表
    val_loss_all = []

    # 训练集准确度的列表
    train_acc_all = []
    # 验证集准确度的列表
    val_acc_all = []
    # 当前时间
    start_time = datetime.now()

    ### 训练过程
    for epoch in range(num_epochs):

        print('-'*30)
        print(f'当前训练轮次={epoch + 1}/{num_epochs}')

        # 每轮的参数
        # 本轮次训练的损失值
        train_loss = 0.0
        # 本轮次训练的准确度
        train_corrects = 0
        # 本轮次验证损失值
        val_loss = 0.0
        # 本轮次验证的准确度
        val_corrects = 0

        # 本轮训练集样本数量
        train_num = 0
        # 本轮验证集样本数量
        val_num = 0

        ##进行每轮训练过程
        for step,(images,labels) in enumerate(train_dataloader):
            ## 取一个‘batch’的数据和标签
            # 数据加载到设备
            images = images.to(device)
            # 标签加载到设备
            labels = labels.to(device)
            # 模型开启训练模式
            model.train()

            # 前向传播：获取分类预测结果,'batch'大小的一个向量，每个标签的预测值[[],[],[]]
            output = model(images)

            # 预测的标签对应的index，获取一行中最大的index索引 ，（在batch行中按行计算每行最大值的索引，即获取预测的标签对于的索引）
            pre_lab = torch.argmax(output,dim=1)

            # 计算batch 的损失值,模型输出与标签计算损失值
            loss = criterion(output,labels)

            # 将梯度初始化为0
            optimizer.zero_grad()

            # 反向传播计算
            loss.backward()

            # 根据反向传播的梯度信息来更新网络的参数，以起到降低loss函数计算值的作用
            optimizer.step()

            #轮次的loss值累加，size(0) 取第一维度的大小（数量）【loss.item() 是该batch的平均损失值】
            train_loss += loss.item()*images.size(0)

            # 轮次的准确度累加,预测正确，则准确度加一
            train_corrects += torch.sum(pre_lab == labels.data)

            # 每批次的数据量累加（size(0) 取第一维度的大小（数量）
            train_num += images.size(0)

            if epoch==0 and step==0 :
                print(output,'前向传播结果')
                print('-'*30)
                print(loss.item(),'轮次的损失值')
                print('-' * 30)
                print(len(images),len(labels),'数据与标签的数量')
                print('-'*30)

        ##进行每轮验证过程
        for step, (images, labels) in enumerate(val_dataloader):
            images = images.to(device)
            labels = labels.to(device)

            # 模型开启验证模式
            model.eval()

            # 前向传播,输入一个batch的图片，得到一个batch中对应的预测值
            output = model(images)

            # 查找batch下每行中最大值对应的索引
            pre_lab = torch.argmax(output,dim=1)

            # 计算每一个batch的loss值
            loss = criterion(output,labels)

            # 验证集对损失值进行累加
            val_loss += loss.item()*images.size(0)
            # 验证集的准确度累加,(对一个batch进行中，预测正确的进行加1)
            val_corrects += torch.sum(pre_lab==labels.data)
            # 验证集训练的样本数量
            val_num += images.size(0)

        ## 计算训练集本轮次的平均loss值
        train_loss_all.append(train_loss / train_num)
        ## 计算训练集本轮次准确率
        train_acc_all.append(train_corrects.double().item() / train_num)

        ## 计算验证集本轮次的平均loss值
        val_loss_all.append(val_loss / val_num)
        ## 计算验证集本轮次的准确率
        val_acc_all.append(val_corrects.double().item() / val_num)

        print(f'{epoch+1}轮次训练，平均loss值={train_loss_all[-1]:.4f},准确度{train_acc_all[-1]:.4f}')
        print(f'{epoch+1}轮次验证，平均loss值={val_loss_all[-1]:.4f},准确度{val_acc_all[-1]:.4f}')


        # 寻找最高准确度
        if val_acc_all[-1]>best_acc:
            # 保存当前的最高准确度
            best_acc = val_acc_all[-1]
            # 保存最高准确度对应的模型参数
            best_model_wts = copy.deepcopy(model.state_dict())
            # 每轮就执行保存最高准确度下的模型参数
            torch.save(best_model_wts, save_path)

        # 计算耗时
        now_time = datetime.now()
        use_seconds = (now_time - start_time).total_seconds()
        time_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f'{epoch+1}轮次训练+验证耗费的时间：{use_seconds//60}分{use_seconds%60}秒',f'此轮结束时间={time_str}')

    ### 训练结束
    # 训练过程每轮的数值进行画图
    train_process = pd.DataFrame(
        data={
            'epoch': range(num_epochs),
            'train_loss_all':train_loss_all,
            'val_loss_all':val_loss_all,
            'train_acc_all':train_acc_all,
            'val_acc_all': val_acc_all,
        }
    )
    return train_process

#训练过程的数据画图出来
def matplot_acc_loss(train_data):
    plt.figure(figsize=(12,4))
    # 1，2，1，  一行两列的第一张图
    plt.subplot(1,2,1)
    plt.plot(train_data['epoch'],train_data.train_loss_all,'ro-',label='train loss')
    plt.plot(train_data['epoch'], train_data.val_loss_all, 'bs-', label='val loss')
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('loss')

    # 1，2，2  一行两列的第二张图
    plt.subplot(1, 2, 2)
    plt.plot(train_data['epoch'], train_data.train_acc_all, 'ro-', label='train acc')
    plt.plot(train_data['epoch'], train_data.val_acc_all, 'bs-', label='val acc')
    plt.legend()
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend()
    plt.show()

# 获取模型对象
def get_model(pth_path:str):
    '''
    :param pth_path:保存模型参数的文件，有该文件就加载参数到模型中，没有模型就使用初始的参数
    :return:
    '''
    model = ModelNet()
    # 有模型参数文件时，就加载模型参数文件
    if os.path.exists(pth_path):
        try:
            load_state = torch.load(pth_path)
            model.load_state_dict(load_state)
        except Exception:
            pass
    return model

## 训练的主进程
def main(num_epoch:int=30):
    '''
    :param num_epoch: 训练的轮次
    :return:
    '''
    # 创建模型文件存放的目录
    create_model_save_path()
    # 模型保存参数的文件
    pth_path = ModelNet.PTH_PATH
    # 实例化模型
    model = get_model(pth_path)
    # 获取训练数据和测试数据
    train_dataloader, val_dataloader = GetDataLoader().train_and_val_dataloader
    # 训练数据，获取每轮训练的数据
    train_data = train_model_process(pth_path,model, train_dataloader, val_dataloader, num_epoch)
    # 画图，每次训练损失值和准确度的变化
    matplot_acc_loss(train_data)

if __name__ == '__main__':
    main(num_epoch=10)
    # train_val_data_process()
